package com.sakura.oj.sandbox.utils;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.command.StatsCmd;
import com.github.dockerjava.api.model.Statistics;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Docker内存监控工具类
 * 用于获取Docker容器的内存使用情况
 */
@Slf4j
@Component
public class DockerMemoryUtils {

    @Resource
    private DockerClient dockerClient;

    private static final int DEFAULT_STATS_TIMEOUT_SECONDS = 3;

    /**
     * 获取Docker容器的内存使用情况（峰值）
     * 
     * @param dockerClient Docker客户端
     * @param containerId 容器ID
     * @param executionTimeMillis 执行时间（毫秒）
     * @return 容器最大内存使用量（KB）
     */
    public static long getMemoryUsage(DockerClient dockerClient, String containerId, long executionTimeMillis) {
        AtomicLong maxMemoryUsage = new AtomicLong(0);
        
        // 设置监控时间
        long timeoutMillis = executionTimeMillis + 500; // 监控时间略大于执行时间
        
        StatsCmd statsCmd = dockerClient.statsCmd(containerId);
        statsCmd.withNoStream(false); // 使用流式获取

        try {
            statsCmd.exec(new MemoryStatsCallback(maxMemoryUsage, false))
                    .awaitCompletion(timeoutMillis, TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            // 如果中断，返回当前最大值
            log.warn("内存监控被中断，返回当前获取的最大值");
        } catch (Exception e) {
            log.error("获取内存使用情况时发生错误", e);
        }

        // 如果没有获取到有效值，尝试单次获取
        if (maxMemoryUsage.get() == 0) {
            return getMemoryUsageSingle(dockerClient, containerId);
        }

        // 转换为KB
        return maxMemoryUsage.get() / 1024;
    }
    
    /**
     * 单次获取容器内存使用情况
     * 
     * @param containerId 容器ID
     * @return 内存使用量（KB）
     */
    public long getMemoryUsageSingle(String containerId) {
        return getMemoryUsageSingle(this.dockerClient, containerId);
    }
    
    /**
     * 单次获取容器内存使用情况
     * 
     * @param dockerClient Docker客户端
     * @param containerId 容器ID
     * @return 内存使用量（KB）
     */
    public static long getMemoryUsageSingle(DockerClient dockerClient, String containerId) {
        AtomicLong memoryUsage = new AtomicLong(0);
        CountDownLatch latch = new CountDownLatch(1);
        
        try {
            StatsCmd statsCmd = dockerClient.statsCmd(containerId);
            statsCmd.withNoStream(true); // 单次获取
            
            statsCmd.exec(new MemoryStatsCallback(memoryUsage, true, latch));
            
            // 等待获取完成或超时
            boolean completed = latch.await(DEFAULT_STATS_TIMEOUT_SECONDS, TimeUnit.SECONDS);
            if (!completed) {
                log.warn("获取容器内存使用情况超时");
            }
        } catch (Exception e) {
            log.error("单次获取内存使用情况时发生错误", e);
        }
        
        return memoryUsage.get() / 1024;
    }

    /**
     * 内存统计回调类
     */
    private static class MemoryStatsCallback extends com.github.dockerjava.api.async.ResultCallback.Adapter<Statistics> {
        private final AtomicLong memoryUsage;
        private final boolean isSingleStat;
        private final CountDownLatch latch;

        public MemoryStatsCallback(AtomicLong memoryUsage, boolean isSingleStat) {
            this(memoryUsage, isSingleStat, null);
        }
        
        public MemoryStatsCallback(AtomicLong memoryUsage, boolean isSingleStat, CountDownLatch latch) {
            this.memoryUsage = memoryUsage;
            this.isSingleStat = isSingleStat;
            this.latch = latch;
        }

        @Override
        public void onNext(Statistics stats) {
            try {
                if (stats != null && stats.getMemoryStats() != null) {
                    // 获取当前内存使用量
                    long currentUsage = stats.getMemoryStats().getUsage();
                    
                    if (isSingleStat) {
                        // 单次统计模式
                        memoryUsage.set(currentUsage);
                        if (latch != null) {
                            latch.countDown();
                        }
                        try {
                            close();
                        } catch (IOException e) {
                            // 忽略关闭异常
                        }
                    } else {
                        // 持续监控模式，更新最大使用量
                        long currentMax = memoryUsage.get();
                        if (currentUsage > currentMax) {
                            memoryUsage.set(currentUsage);
                        }
                    }
                    
                    log.debug("容器内存使用: {} 字节", currentUsage);
                }
            } catch (Exception e) {
                log.warn("处理内存统计数据时出错", e);
                if (latch != null) {
                    latch.countDown();
                }
            }
            super.onNext(stats);
        }

        @Override
        public void onError(Throwable throwable) {
            // 忽略管道关闭异常
            if (!(throwable instanceof IOException && throwable.getMessage() != null &&
                (throwable.getMessage().contains("管道已") || 
                 throwable.getMessage().contains("pipe") || 
                 throwable.getMessage().contains("Pipe") || 
                 throwable.getMessage().contains("closed")))) {
                log.error("获取内存统计时出错", throwable);
            }
            
            if (latch != null) {
                latch.countDown();
            }
        }
        
        @Override
        public void onComplete() {
            if (latch != null) {
                latch.countDown();
            }
            super.onComplete();
        }
    }
} 